import logging
import sys
import time
from binascii import hexlify
import zmq
import MDP
from zhelpers import dump


# 定义一个服务
class Service(object):
    # 服务名称
    name = None
    # 客户端请求队列
    requests = None
    # 正在等待的worker队列
    waiting = None

    def __init__(self, name):
        self.name = name
        self.requests = []
        self.waiting = []


# 定义一个worker，状态为空闲或占用
class Worker(object):
    # worker的标识
    identity = None
    # 地址帧
    address = None
    # 所属服务
    service = None
    # 过期时间
    expiry = None

    def __init__(self, identity, address, lifetime):
        self.identity = identity
        self.address = address
        self.expiry = time.time() + 1e-3*lifetime


# 定义一个代理
class MajorDomoBroker(object):
    INTERNAL_SERVICE_PREFIX = b"mmi."
    HEARTBEAT_LIVENESS = 3
    HEARTBEAT_INTERVAL = 2500
    HEARTBEAT_EXPIRY = HEARTBEAT_INTERVAL * HEARTBEAT_LIVENESS

    # 上下文
    ctx = None
    # 用于连接client和worker的套接字
    socket = None
    poller = None

    # 发送心跳的时间
    heartbeat_at = None
    # 已知服务的哈希表
    services = None
    # 已知worker的哈希表
    workers = None
    # 正在等待的worker队列
    waiting = None

    # 使用标准输出打印活动信息
    verbose = False

    def __init__(self, verbose=False):
        self.verbose = verbose
        self.services = {}
        self.workers = {}
        self.waiting = []
        self.heartbeat_at = time.time() + 1e-3*self.HEARTBEAT_INTERVAL
        self.ctx = zmq.Context()
        self.socket = self.ctx.socket(zmq.ROUTER)
        self.socket.linger = 0
        self.poller = zmq.Poller()
        self.poller.register(self.socket, zmq.POLLIN)
        logging.basicConfig(format="%(asctime)s %(message)s", datefmt="%Y-%m-%d %H:%M:%S", level=logging.INFO)

    def bind(self, endpoint):
        self.socket.bind(endpoint)
        logging.info("I: MDP broker/0.1.1 is active at %s" % endpoint)

    # 代理使用的方法
    def mediate(self):
        while True:
            try:
                items = self.poller.poll(self.HEARTBEAT_INTERVAL)
            except KeyboardInterrupt:
                break
            if items:
                msg = self.socket.recv_multipart()
                if self.verbose:
                    logging.info("I: received message:")
                    dump(msg)
                sender = msg.pop(0)
                empty = msg.pop(0)
                assert empty == b""
                header = msg.pop(0)

                if MDP.C_CLIENT == header:
                    self.process_client(sender, msg)
                elif MDP.W_WORKER == header:
                    self.process_worker(sender, msg)
                else:
                    logging.error("E: invalid message:")
                    dump(msg)

            self.purge_workers()
            self.send_heartbeats()

    def destroy(self):
        while self.workers:
            self.delete_worker(self.workers.values()[0], True)
        self.ctx.destroy(0)

    def purge_workers(self):
        while self.waiting:
            w = self.waiting[0]
            curtime = time.time()
            if w.expiry < curtime:
                logging.info("I: expiry: %s curtime: %s deleting expired worker: %s" % (w.expiry, curtime, w.identity))
                self.delete_worker(w, False)
                self.waiting.pop(0)
            else:
                break

    def send_heartbeats(self):
        if time.time() > self.heartbeat_at:
            for worker in self.waiting:
                self.send_to_worker(worker, MDP.W_HEARTBEAT, None, None)
            self.heartbeat_at = time.time() + 1e-3*self.HEARTBEAT_INTERVAL

    # 服务使用的方法
    def require_service(self, name):
        assert name is not None
        service = self.services.get(name)
        if service is None:
            service = Service(name)
            self.services[name] = service
        return service

    # 可能时，分发请求给等待中的worker
    def dispatch(self, service, msg):
        assert service is not None
        if msg is not None:
            service.requests.append(msg)
        self.purge_workers()
        while service.waiting and service.requests:
            msg = service.requests.pop(0)
            worker = service.waiting.pop(0)
            self.waiting.remove(worker)
            self.send_to_worker(worker, MDP.W_REQUEST, None, msg)

    # 使用8/MMI协定处理内部服务
    def service_internal(self, service, msg):
        returncode = b"501"
        if b"mmi.service" == service:
            name = msg[-1]
            returncode = b"200" if name in self.services else b"404"
        msg[-1] = returncode

        msg = msg[:2] + [MDP.C_CLIENT, service] + msg[2:]
        self.socket.send_multipart(msg)

    # worker使用的方法
    def require_worker(self, address):
        assert address is not None
        identity = hexlify(address)
        worker = self.workers.get(identity)
        if worker is None:
            worker = Worker(identity, address, self.HEARTBEAT_EXPIRY)
            self.workers[identity] = worker
            if self.verbose:
                logging.info("I: registering new worker: %s" % identity)
        return worker

    def delete_worker(self, worker, disconnect):
        assert worker is not None
        if disconnect:
            self.send_to_worker(worker, MDP.W_DISCONNECT, None, None)
        if worker.service is not None:
            worker.service.waiting.remove(worker)
        self.workers.pop(worker.identity)

    def process_worker(self, sender, msg):
        assert len(msg) >= 1
        command = msg.pop(0)
        worker_ready = hexlify(sender) in self.workers
        worker = self.require_worker(sender)
        if MDP.W_READY == command:
            assert len(msg) >= 1
            service = msg.pop(0)
            if worker_ready or service.startswith(self.INTERNAL_SERVICE_PREFIX):
                self.delete_worker(worker, True)
            else:
                worker.service = self.require_service(service)
                self.worker_waiting(worker)
        elif MDP.W_REPLY == command:
            if worker_ready:
                client = msg.pop(0)
                empty = msg.pop(0)
                msg = [client, b'', MDP.C_CLIENT, worker.service.name] + msg
                self.socket.send_multipart(msg)
                self.worker_waiting(worker)
            else:
                self.delete_worker(worker, True)
        elif MDP.W_HEARTBEAT == command:
            if worker_ready:
                worker.expiry = time.time() + 1e-3 * self.HEARTBEAT_EXPIRY
            else:
                self.delete_worker(worker, True)
        elif MDP.W_DISCONNECT == command:
            self.delete_worker(worker, False)
        else:
            logging.error("E: invalid message:")
            dump(msg)

    def send_to_worker(self, worker, command, option, msg=None):
        if msg is None:
            msg = []
        elif not isinstance(msg, list):
            msg = [msg]

        if option is not None:
            msg = [option] + msg

        msg = [worker.address, b'', MDP.W_WORKER, command] + msg

        if self.verbose:
            logging.info("I: sending %r to worker" % command)
            dump(msg)

        self.socket.send_multipart(msg)

    def worker_waiting(self, worker):
        self.waiting.append(worker)
        worker.service.waiting.append(worker)
        worker.expiry = time.time() + 1e-3*self.HEARTBEAT_EXPIRY
        self.dispatch(worker.service, None)

    # 客户端使用的方法
    def process_client(self, sender, msg):
        assert len(msg) >= 2
        service = msg.pop(0)
        msg = [sender, b''] + msg
        if service.startswith(self.INTERNAL_SERVICE_PREFIX):
            self.service_internal(service, msg)
        else:
            self.dispatch(self.require_service(service), msg)


def main():
    verbose = '-v' in sys.argv
    broker = MajorDomoBroker(verbose)
    broker.bind("tcp://*:5555")
    broker.mediate()


if __name__ == "__main__":
    main()
